package com.baidu.qianfan.service;

import com.baidu.qianfan.entity.ERNIE4ResponseBody;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.web.bind.annotation.Mapping;

import java.io.IOException;
import java.time.Duration;
import java.util.Collections;

/**
 * ERNIE-4.0-8K的调用方法
 * 百度千帆官方文档是：https://cloud.baidu.com/doc/WENXINWORKSHOP/s/clntwmv7t
 */
@Service
public class ERNIE4 {

    @Value("${baidu.qianfan.API_KEY}")
    private String API_KEY;
    @Value("${baidu.qianfan.SECRET_KEY}")
    private String SECRET_KEY;

    // 发起请求，调用接口
    public  ERNIE4ResponseBody extractedERNIE4(String text) throws Exception {
        String accessToken = getAccessToken();
        String url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + accessToken;
        String payload = new Gson().toJson(buildPayload(text));

        MediaType JSON = MediaType.parse("application/json; charset=utf-8");
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(Duration.ofSeconds(120)) // 连接超时时间 秒
                .writeTimeout(Duration.ofSeconds(120))   // 写入超时时间
                .readTimeout(Duration.ofSeconds(120))    // 读取超时时间
                .build();
        RequestBody body = RequestBody.create(payload, JSON);
        Request request = new Request.Builder()
                .url(url)
                .post(body)
                .addHeader("Content-Type", "application/json")
                .build();
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) throw new IOException("意外代码：" + response);
            // 读取响应体为字符串
            String responseBody = response.body().string();
            // 使用Gson将字符串转换为Java对象
            Gson gson = new Gson();
            ERNIE4ResponseBody ernie4ResponseBody = gson.fromJson(responseBody, ERNIE4ResponseBody.class);
            return ernie4ResponseBody;
        }
    }

    // 构建加载请求参数
    private JsonObject buildPayload(String text) {
        JsonObject messages = new JsonObject();
        messages.addProperty("role", "user");
        messages.addProperty("content", text);

        JsonObject payload = new JsonObject();
        payload.add("messages", new Gson().toJsonTree(Collections.singletonList(messages)));
        payload.addProperty("temperature", 0.95);
        payload.addProperty("top_p", 0.8);
        payload.addProperty("penalty_score", 1);
        payload.addProperty("enable_system_memory", false);
        payload.addProperty("disable_search", false);
        payload.addProperty("enable_citation", false);

        return payload;
    }

    // 获取token
    private String getAccessToken() {
        OkHttpClient client = new OkHttpClient();
        String url = "https://aip.baidubce.com/oauth/2.0/token";

        RequestBody formBody = new FormBody.Builder()
                .add("grant_type", "client_credentials")
                .add("client_id", API_KEY)
                .add("client_secret", SECRET_KEY)
                .build();

        Request request = new Request.Builder()
                .url(url)
                .post(formBody)
                .build();

        try (Response response = client.newCall(request).execute()) {
            return new Gson().fromJson(response.body().string(), JsonObject.class).get("access_token").getAsString();
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }
}